# include <yolov8_yolo11_yolo12_det.h>

class Logger : public ILogger
{
	void log(Severity severity, const char* msg)  noexcept
	{
		// suppress info-level messages
		if (severity != Severity::kINFO)
			std::cout << msg << std::endl;
	}
} gLogger;

YOLOv81112TRTDetector::~YOLOv81112TRTDetector() {
	// ???????????????
	cudaStreamSynchronize(stream);
	cudaStreamDestroy(stream);
	if (!buffers[0]) {
		delete[] buffers;
	}
	std::cout << "destory instance done!" << std::endl;
}

void YOLOv81112TRTDetector::initConfig(std::string enginepath, float conf, float scored) {
	std::ifstream file(enginepath, std::ios::binary);
	char* trtModelStream = NULL;
	int size = 0;
	if (file.good()) {
		file.seekg(0, file.end);
		size = file.tellg();
		file.seekg(0, file.beg);
		trtModelStream = new char[size];
		assert(trtModelStream);
		file.read(trtModelStream, size);
		file.close();
	}
	this->runtime = createInferRuntime(gLogger);
	assert(this->runtime != nullptr);
	this->engine = this->runtime->deserializeCudaEngine(trtModelStream, size);
	assert(this->engine != nullptr);
	this->context = engine->createExecutionContext();
	assert(this->context != nullptr);
	delete[] trtModelStream;

	// ???????????????
	auto inputDims = this->engine->getTensorShape("images");
	auto outDims = this->engine->getTensorShape("output0");

	// ?????????????
	this->input_h = inputDims.d[2];
	this->input_w = inputDims.d[3];
	printf("inputH : %d, inputW: %d \n", this->input_h, this->input_w);

	// ????????????
	this->output_h = outDims.d[1];
	this->output_w = outDims.d[2];
	std::cout << "out data format: " << this->output_h << "x" << this->output_w << std::endl;

	// ????GPU???????/?????????
	cudaMalloc(&buffers[0], this->input_h * input_w * 3 * sizeof(float));
	cudaMalloc(&buffers[1], this->output_h * this->output_w * sizeof(float));

	// ??????????????
	this->prob.resize(this->output_h * this->output_w);

	// ????cuda??
	cudaStreamCreate(&stream);
}

void YOLOv81112TRTDetector::detect(cv::Mat& frame, std::vector<DetectResult>& results) {
	int64 start = cv::getTickCount();
	// ???????? - ?????????
	int w = frame.cols;
	int h = frame.rows;
	int _max = std::max(h, w);

	// 图象预处理 - 格式化操作
	cv::cuda::GpuMat gpu_frame;
	gpu_frame.upload(frame);

	// cv::Mat image = cv::Mat::zeros(cv::Size(_max, _max), CV_8UC3);
	// cv::Rect roi(0, 0, w, h);
	// frame.copyTo(image(roi));

	cv::cuda::GpuMat gpu_image = cv::cuda::GpuMat(cv::Size(_max, _max), CV_8UC3);
	cv::Rect roi(0, 0, w, h);
	gpu_frame.copyTo(gpu_image(roi));

	// float x_factor = image.cols / static_cast<float>(this->input_w);
	// float y_factor = image.rows / static_cast<float>(this->input_h);

	float x_factor = gpu_image.cols / static_cast<float>(input_w);
	float y_factor = gpu_image.rows / static_cast<float>(input_h);

	// 图像缩放
	// 归一化（1/255.0）
	// 通道转换（BGR→RGB）
	// 转换为 CHW 格式的 blob 
	// cv::Mat blob = cv::dnn::blobFromImage(image, 1 / 255.0, cv::Size(this->input_w, this->input_h), cv::Scalar(0, 0, 0), true, false);
	// cudaMemcpyAsync(buffers[0], blob.ptr<float>(), 3 * this->input_h * this->input_w * sizeof(float), cudaMemcpyHostToDevice, stream);

	// 图像缩放
	cv::cuda::GpuMat resized;
	cv::cuda::resize(gpu_image, resized, cv::Size(input_w, input_h), 0, 0, cv::INTER_LINEAR);
	// 归一化（1/255.0）
	cv::cuda::GpuMat flt_image;
	resized.convertTo(flt_image, CV_32FC3, 1.f / 255.f);
	// 通道转换（BGR→RGB）
	cv::cuda::GpuMat rgb;
	cv::cuda::cvtColor(flt_image, rgb, cv::COLOR_BGR2RGB);
	// HWC 转换为 CHW 格式
	// 将 OpenCV 默认的 HWC 格式（高度、宽度、通道数）转换为深度学习模型常用的 CHW 格式（通道数、高度、宽度），
	// 同时直接将数据写入模型的 GPU 输入缓冲区，避免了额外的内存复制，提高了预处理效率。

	// 这部分是为每个通道创建一个GpuMat，并直接将其数据指针绑定到模型输入缓冲区（buffers[0]）的对应位置：
	// buffers[0]是 TensorRT 的 GPU 输入缓冲区，用于存放模型输入数据（CHW 格式）。
	// 循环中通过指针偏移（i * input_w * input_h）为 R、G、B 三个通道分别分配缓冲区中的位置：
	// R 通道数据将写入buffers[0]的起始位置
	// G 通道数据将写入buffers[0]偏移input_w * input_h的位置
	// B 通道数据将写入buffers[0]偏移2 * input_w * input_h的位置
	std::vector<cv::cuda::GpuMat> rgb_out;
	for (size_t i = 0; i < 3; ++i)
	{
		rgb_out.emplace_back(cv::cuda::GpuMat(cv::Size(input_w, input_h), CV_32FC1, (float*)buffers[0] + i * input_w * input_h));
	}
	// 这是 GPU 版的通道分割函数，功能是将输入的三通道图像rgb（HWC 格式）分割为三个单通道图像，并存储到rgb_out向量中：
	// rgb：输入的三通道图像（GPU 上的CV_32FC3类型，RGB 格式，HWC 布局）。
	// rgb_out：输出的单通道向量，分割后rgb_out[0]为 R 通道，rgb_out[1]为 G 通道，rgb_out[2]为 B 通道。
	// 由于rgb_out的每个元素已绑定到buffers[0]的对应位置，分割后的数据会直接写入模型输入缓冲区，无需额外复制。
	// opencv HWC order -> CHW order
	cv::cuda::split(rgb, rgb_out);

	// 释放该对象所占用的 GPU 内存
	gpu_frame.release();
	resized.release();
	flt_image.release();
	rgb.release();
	gpu_image.release();

	
	// 推理
	context->executeV2(buffers);

	// GPU显存到内存
	cudaMemcpyAsync(prob.data(), buffers[1], this->output_h * this->output_w * sizeof(float), cudaMemcpyDeviceToHost, stream);

	// 后处理, 1x25200x85  
	std::vector<cv::Rect> boxes;
	std::vector<int> classIds;
	std::vector<float> confidences;
	cv::Mat dout(84, 8400, CV_32F, (float*)prob.data());
	cv::Mat det_output = dout.t(); // 8400x84
	for (int i = 0; i < det_output.rows; i++) {
		cv::Mat classes_scores = det_output.row(i).colRange(4, 84);
		cv::Point classIdPoint;
		double score;
		minMaxLoc(classes_scores, 0, &score, 0, &classIdPoint);

		// 置信度 0～1之间
		if (score > this->score_threshold)
		{
			float cx = det_output.at<float>(i, 0);
			float cy = det_output.at<float>(i, 1);
			float ow = det_output.at<float>(i, 2);
			float oh = det_output.at<float>(i, 3);
			int x = static_cast<int>((cx - 0.5 * ow) * x_factor);
			int y = static_cast<int>((cy - 0.5 * oh) * y_factor);
			int width = static_cast<int>(ow * x_factor);
			int height = static_cast<int>(oh * y_factor);
			cv::Rect box;
			box.x = x;
			box.y = y;
			box.width = width;
			box.height = height;

			boxes.push_back(box);
			classIds.push_back(classIdPoint.x);
			confidences.push_back(score);
		}
	}

	// NMS
	std::vector<int> indexes;
	cv::dnn::NMSBoxes(boxes, confidences, 0.25, 0.45, indexes);
	for (size_t i = 0; i < indexes.size(); i++) {
		DetectResult dr;
		int index = indexes[i];
		int idx = classIds[index];
		dr.box = boxes[index];
		dr.classId = idx;
		dr.score = confidences[index];
		// cv::rectangle(frame, boxes[index], cv::Scalar(0, 0, 255), 1, 8);
		// cv::rectangle(frame, cv::Point(boxes[index].tl().x, boxes[index].tl().y - 20),
		// 	cv::Point(boxes[index].br().x, boxes[index].tl().y), cv::Scalar(0, 255, 255), -1);
		results.push_back(dr);
	}

	// ????FPS render it
	float t = (cv::getTickCount() - start) / static_cast<float>(cv::getTickFrequency());
	// putText(frame, cv::format("FPS: %.2f", 1.0 / t), cv::Point(20, 40), cv::FONT_HERSHEY_PLAIN, 2.0, cv::Scalar(0, 255, 0), 2, 8);
}
